#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include "scalar.h"
#include "optim.h"
/**
 *  y = sqrt(x)
 *  Loss(y) = (y^2 - x)^2
 */

 ScalarPtr loss(ScalarPtr y, ScalarPtr x) 
 {
    return Pow(Sub(Pow(y, NewScalar(2, false)), x), NewScalar(2, false));
 }

 ScalarPtr Sqrt(ScalarPtr p)
 {
    ScalarPtr t = NewScalar(1e-1, true);
    SGDPtr optim = NewSGD(t, 1e-3);
    while (true) {
        ScalarPtr e = loss(t, p);
        // printf("t = %.7f loss = %.7f\n", Scalar2Double(t), Scalar2Double(e));
        if (! GT(e, 1e-8)) {
            break;
        }
        ZeroGrad(optim);
        Backward(e);
        Step(optim);
    }
    return t;
 }

int main(int argc, char *argv[])
{
    assert(argc == 2);
    ScalarPtr x = NewScalar(atof(argv[1]), false);
    ScalarPtr y = Sqrt(x);
    printf("%.3f\n", Scalar2Double(y));
    return 0;
}